#!/usr/bin/env python
import numpy as np
import itk
import scipy.ndimage as ndi
import sys
import os

def lungMaskExtraction(fileName):
  reader = itk.ImageFileReader.IF3.New()
  reader.SetFileName(fileName)
  writer = itk.ImageFileWriter.IUS3.New()
  tmp = os.path.splitext(fileName) 
  foutput = tmp[0] + "_lungMask" + tmp[1]
  writer.SetFileName(foutput)
  btf = itk.BinaryThresholdImageFilter.IF3IUS3.New()
  btf.SetInput(reader.GetOutput())
  btf.SetLowerThreshold(-900)
  btf.SetUpperThreshold(-200)
  ccf = itk.ConnectedComponentImageFilter.IUS3IUL3.New()
  ccf.SetInput(btf.GetOutput())
  rcf = itk.RelabelComponentImageFilter.IUL3IUS3.New()
  rcf.SetInput(ccf.GetOutput())
  btf2 = itk.BinaryThresholdImageFilter.IUS3IUS3.New()
  btf2.SetInput(rcf.GetOutput())
  btf2.SetLowerThreshold(1)
  btf2.SetUpperThreshold(1)
  btf2.Update()
  img = itk.PyBuffer.IUS3.GetArrayFromImage(btf2.GetOutput())
  img = np.where(img>0,1,0).astype(np.uint8)

  newImg = np.zeros(img.shape,np.uint8)
  for i in xrange(newImg.shape[0]):
      ns = ndi.binary_fill_holes(img[i,:,:])
      newImg[i,:,:] = ns

  lungImg = itk.PyBuffer.IUS3.GetImageFromArray(newImg)
  writer.SetInput(lungImg)
  writer.Update()
  return foutput

def main():
  if (len(sys.argv) < 2):
     print ("Not enough arguments")
     sys.exit(1) 

  print sys.argv[1]
  inputFile = sys.argv[1]

  lungMaskExtraction(inputFile)

if __name__ == '__main__':
    main()
 
